import hashlib
import torch
import torch.nn as nn
import feature_extractor
from CoSMAE import CoSMAE
import torchvision.models.vision_transformer as vision_transformer


def get_backbone(image_size, number_of_channels, patch_size):
    model = vision_transformer.VisionTransformer(image_size=image_size,num_classes=17,patch_size=patch_size,hidden_dim=768,num_layers=12,num_heads=12,mlp_dim=3072) 
    model.conv_proj= nn.Conv2d(in_channels=number_of_channels, out_channels=768, kernel_size=patch_size, stride=patch_size) 
      
    return model



def get_ssl_model(model, ssl_flag, mr=0.75):
    if ssl_flag == 'CoSMAE':
        ssl_model = CoSMAE(model, mr)
        return ssl_model

    
def get_pretrained_weights(model, path, device):
    print('load weights from {}'.format(path))
    pretrained_state_dict = torch.load('weights/' + path, map_location=torch.device(device))
    model_state_dict = model.state_dict()  # Your model's state dictionary
    pre_filter = {k: v for k, v in pretrained_state_dict.items() if k.split('.')[0] != 'Projector'}
    print(len(pretrained_state_dict.keys()))
    print(len(pre_filter.keys()))
    filtered_state_dict = {k: v for k, v in pre_filter.items() if k in model_state_dict}
    print(len(filtered_state_dict.keys()))
    #filtered_state_dict = {k: v for k, v in pretrained_state_dict.items() if k in model_state_dict}

    # Update the model's state dictionary with the filtered one
    model_state_dict.update(filtered_state_dict)

    # Load the updated state dictionary into your model
    model.load_state_dict(model_state_dict)


def get_old_encoder(model, device, patch_size):
    backbone_copy = get_backbone(224,3, patch_size)#(256, 3)
    ssl_model_copy = get_ssl_model(backbone_copy, 'CoSMAE')
    #model_copy = get_ssl_model(backbone_copy, 'MAE')
    #model_copy = type(model)()  # create a new instance of the same class
    ssl_model_copy.backbone.load_state_dict(model.backbone.state_dict())#backbone_copy.load_state_dict(model.backbone.state_dict())
    ssl_model_copy.backbone.to(device)

    old_encoder = feature_extractor.FeatureExtractor(ssl_model_copy.backbone,'encoder')
    for param in old_encoder.parameters():
        param.requires_grad = False
    old_encoder.to(device)
    return old_encoder


def map_dataset_to_file(name, split):
    
    if name == 'US3D':
        if split == 'train':
            path = './paths/JAX_256_rgb_train_30.txt'
        else:
            path = './paths/JAX_256_rgb_{}.txt'.format(split)
    elif name == 'UAVID':
        path = './paths/UAVID_256_{}.txt'.format(split)
    elif name == 'Potsdam':
        path = './paths/Potsdam_256_rgb_{}.txt'.format(split)
    elif name == 'TreeSatAI':
        if split == 'train':
            path = './paths/treesatai_train_15000_0.txt'
        else:
            path = './paths/treesatai_{}.txt'.format(split)
    elif name == 'UCMerced':
        path = './paths/UCMerced_1000_{}_0.txt'.format(split)
         
    print(path)
    return path
    

def file_to_list(path):
    with open(path, 'r') as f:
        paths_list = f.read().split('\n')[:-1]
    return paths_list

def get_model_checksum(model):
    md5 = hashlib.md5()
    for param in model.parameters():
        md5.update(param.data.cpu().numpy().tobytes())
    return md5.hexdigest()